# -*- coding: utf-8 -*-
"""
@Time ： 2020-11-19 9:30
@Auth ： lixin
@Description：数值型数据回归算法集合训练类

"""
from sklearn import linear_model

from algo.Algo_interface import Algo_interface
from sklearn import linear_model
from sklearn import metrics
from sklearn.ensemble import RandomForestClassifier
from sklearn import svm
import api as runs
import numpy as np
import lightgbm as lgb
from sklearn import neural_network
import xgboost as xgb
import json

class DigitalTypeRegressor(Algo_interface):
    def __init__(self, model_type, model_name, model_params):
        self.task_type = model_type
        self.model_name = model_name
        self.model_params = model_params
        self.model = None
        self.build_model()
        # return self.model

    def set_model(self, model):
        self.model = model
        return 1

    def get_model(self):
        return self.model

    def build_model(self):
        if self.model_name == '线性回归':
            self.model = linear_model.LinearRegression(**self.model_params)
        elif self.model_name == 'lgbm':
            self.model = lgb.LGBMRegressor(**self.model_params)
        elif self.model_name == '神经网络':
            self.model = neural_network.MLPRegressor(**self.model_params)
        elif self.model_name == 'svr':
            self.model = svm.SVR(**self.model_params)
        elif self.model_name == 'xgboost':
            self.model = xgb.XGBRegressor(**self.model_params)
    def train(self,data):
        x_train, y_train=data
        self.model.fit(x_train, y_train)
        return 1
    def predict(self,data):
        y_pred = self.model.predict(data)
        return y_pred
